# pip install numpy xarray netCDF4 matplotlib plotly
# pip install "dask[complete]" -q
# pip install geopy ipywidgets folium -q
# pip install pycountry -q
# pip install shapely
# pip install -U kaleido
# !jupyter nbextension enable --py widgetsnbextension
# !jupyter labextension install @jupyter-widgets/jupyterlab-manager

from ipywidgets import Layout, Dropdown, widgets
from IPython.display import display, clear_output, IFrame
from functools import partial
import datetime
import numpy as np
import modules.n1_utilities as uti
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)

country_list = uti.read_json_to_sorted_dict('countries.json')
months = uti.read_json_to_dict('months.json')
timescales = uti.read_json_to_dict('timescales.json')
subset_area = None
index = None
bounding_box = (None, None, None, None)
active_btn = None

selected = {
    "country": None,
    "adm1_subarea": None,
    "adm2_subarea": None,
    "timescale": None,
    "month": None,
    "year": None,
    "year_range": None
}

placeholders = {
    "country": "no country selected...",
    "adm1_subarea": "no adm1 subarea selected...",
    "adm2_subarea": "no adm2 subarea selected...",
    "timescale": "no timescale selected...",
    "month": "no month selected...",
    "year": "no year selected..."
}
uti.save_selection(placeholders)

# Custom style and layout for descriptions and dropdowns
style = {'description_width': '150px'}
dropdown_layout = Layout(width='400px', display='flex', justify_content='flex-end')
range_layout = Layout(width='400px')
btn_layout = Layout(width='400px')


# Dropdown for countries
country_names = [country['name'] for country in country_list]
country_selector = widgets.Dropdown(
    options=[placeholders['country']] + country_names,
    description='Select a country:',
    style=style,
    layout=dropdown_layout
)

# Dropdown for subareas, initially empty
adm1_subarea_selector = widgets.Dropdown(
    options=[placeholders['adm1_subarea']],
    description='a subarea of first level:',
    style=style,
    layout=dropdown_layout
)

adm2_subarea_selector = widgets.Dropdown(
    options=[placeholders['adm2_subarea']],
    description='or of second level:',
    style=style,
    layout=dropdown_layout
)

# Dropdown for timescales
timescale_selector = widgets.Dropdown(
    options=[placeholders['timescale']] + list(timescales.keys()),
    description='Select a timescale:',
    style=style,
    layout=dropdown_layout
)

# Dropdown for months
month_selector = widgets.Dropdown(
    options=[placeholders['month']] + list(months.keys()),
    description='Select a month:',
    style=style,
    layout=dropdown_layout
)

# Dropdown for years
current_year = datetime.datetime.now().year
years_options = [str(year) for year in range(1940, current_year + 1)]

year_selector = widgets.Dropdown(
    options=[placeholders['year']] + years_options,
    description='Select a year:',
    disabled=False,
    style=style,
    layout=dropdown_layout
)

# SelectionRangeSlider for years
year_range_selector = widgets.SelectionRangeSlider(
    options=years_options,
    index=(len(years_options) - 1, len(years_options) - 1),  # Start and end at the last
    description='Select the year range:',
    disabled=False,
    style=style,
    layout=range_layout
)

selectors = {
    "country" : country_selector,
    "adm1_subarea": adm1_subarea_selector,
    "adm2_subarea": adm2_subarea_selector,
    "timescale": timescale_selector,
    "month": month_selector,
    "year": year_selector,
    "year_range": year_range_selector    
}


month_widgets_btn = widgets.Button(
    description='Get data',
    disabled=False,
    button_style='info', # 'success', 'info', 'warning', 'danger' or ''
    tooltip='Click me',
    icon='filter', # (FontAwesome names without the `fa-` prefix)month
    layout=btn_layout
)
month_widgets_btn.custom_name='month_widgets_btn'


year_widgets_btn = widgets.Button(
    description='Get data',
    disabled=False,
    button_style='info', # 'success', 'info', 'warning', 'danger' or ''
    tooltip='Click me',
    icon='filter', # (FontAwesome names without the `fa-` prefix)
    layout=btn_layout
)
year_widgets_btn.custom_name='year_widgets_btn'

year_range_widgets_btn = widgets.Button(
    description='Get data',
    disabled=False,
    button_style='info', # 'success', 'info', 'warning', 'danger' or ''
    tooltip='Click me',
    icon='filter', # (FontAwesome names without the `fa-` prefix)
    layout=btn_layout
)
year_range_widgets_btn.custom_name='year_range_widgets_btn'

# Output area for display updates
output_area = widgets.Output()

def setup_observers():
    """
    Sets up observers for UI widgets to handle interactions and updates dynamically in a graphical user interface.
    This function ensures that observers are only set once using a function attribute to track whether observers have
    already been established, enhancing efficiency and preventing multiple bindings to the same event.

    Observer is attached to widgets for country selection. This observer triggers specific functions when the 'value' property 
    of the widgets changes, facilitating responsive updates to the user interface
    based on user interactions.

    Notes:
    - This function uses a custom attribute `observers_set` on itself to ensure observers are set only once.
    """
    if not hasattr(setup_observers, 'observers_set'):      
            # When 'value' changes, update_subareas function will be called to update the dropdown menus
            # Create a partial function that includes the additional parameters
            country_selector.observe(partial(uti.update_subareas, 
                                         country_list=country_list, 
                                         placeholders=placeholders,
                                         adm1_subarea_selector=adm1_subarea_selector, 
                                         adm2_subarea_selector=adm2_subarea_selector), 'value')
            # Set a flag to indicate observers are set
            setup_observers.observers_set = True


            
# Function to update and get data
def update_and_get_data(btn_name):
    global selected, placeholders, output_area, subset_data, index, bounding_box, active_btn
    map_display = None
    active_btn = btn_name
    uti.month_year_interaction(btn_name, month_selector, year_selector, selected, placeholders)
    if uti.validate_selections(btn_name, selected, selectors, placeholders, output_area):
        with output_area:
            output_area.clear_output(wait=True)
            coordinates = uti.get_boundaries(selected, country_list, placeholders)
            # print(coordinates)
            bounding_box = uti.calculate_bounding_box(coordinates)
            # print(bounding_box)
            
                        
            # sample_coordinates = coordinates[:3] # Showing first 3 coordinates for brevity            
            # print('Original Coordinates Sample: ', sample_coordinates)  
            # print('Bounding Box: ', bounding_box)
                        
            # Fetching data using the bounding box
            subset_data = uti.get_xarray_data(btn_name, bounding_box, selectors, placeholders, months, timescales)
            index = f"SPEI{timescales[selectors['timescale'].value]}"
            adm_level, selected_area = uti.get_adm_level_and_area_name(selected, placeholders)
            timescale = selected['timescale']
            time_period = uti.get_period_of_time(btn_name, selected, placeholders)
                
            print(f"SPEI subset data uploaded for {selected_area}, administrative level {adm_level}, timescale {timescale}, period {time_period}")
            zoom_start = 4
            if adm_level == 'ADM1' or adm_level == 'ADM2':
                zoom_start = 8  
            map_display = uti.display_map(bounding_box, zoom_start)
            map_iframe = uti.display_map_in_iframe(map_display)
            display(map_iframe)

            
# Set up widget interaction
def on_button_clicked(btn):
    update_and_get_data(btn.custom_name)


# Setup observers
setup_observers()

# Update existing selectors
previous_selection = uti.read_json_to_dict('selection.json')

# Set up widgets with previous settings
country_selector.value = previous_selection.get('country', placeholders['country'])
adm1_subarea_selector.value = previous_selection.get('adm1_subarea', placeholders['adm1_subarea'])
adm2_subarea_selector.value = previous_selection.get('adm2_subarea', placeholders['adm2_subarea'])
timescale_selector.value = previous_selection.get('timescale', placeholders['timescale'])
month_selector.value = previous_selection.get('month', placeholders['month'])
month_widgets_btn.on_click(on_button_clicked)

# Display widgets
display(country_selector, adm1_subarea_selector, adm2_subarea_selector, timescale_selector, month_selector, month_widgets_btn, output_area)

uti.display_data_details(active_btn, selected, subset_data[index])

Country:  Italy
ADM1 subarea:  no adm1 subarea selected...
ADM2 subarea:  no adm2 subarea selected...
Month:  March
Timescale:  3 months 

Time values in the subset: 84
Latitude values in the subset: 47
Longitude values in the subset: 48 

Data sample:  [[ 0.10656517 -0.03892938 -0.21370953 -0.26820409 -0.31034104]
 [ 0.20403971  0.12018371 -0.05364322 -0.16029305 -0.19767122]
 [ 0.09961821  0.06639799  0.11098377  0.02562103  0.10002738]
 [-0.0285368  -0.03788812  0.03719397  0.16075343  0.21798218]
 [-0.13031167 -0.14410005 -0.00440113  0.11054687  0.14219684]]

processed_subset, change_summary = uti.process_datarray(subset_data[index])
print(processed_subset, '\n')
print('Change summary:')
for key, val in change_summary.items():
    print(key, val)

<xarray.DataArray 'SPEI3' (time: 84, lat: 47, lon: 48)>
dask.array<where, shape=(84, 47, 48), dtype=float64, chunksize=(1, 47, 48), chunktype=numpy.ndarray>
Coordinates:
  * time     (time) datetime64[ns] 1940-03-01T06:00:00 ... 2023-03-01T06:00:00
  * lon      (lon) float64 6.75 7.0 7.25 7.5 7.75 ... 17.5 17.75 18.0 18.25 18.5
  * lat      (lat) float64 35.5 35.75 36.0 36.25 36.5 ... 46.25 46.5 46.75 47.0
Attributes:
    long_name:  Standardized Drought Index (SPEI3)
    units:      - 

Change summary:
invalid_values_replaced 86268
invalid_ratio 0.45523049645390073
duplicates_removed 0
cftime_conversions 0

stat_values = uti.compute_stats(processed_subset)

uti.create_scatterplot(stat_values, timescales, selected, placeholders)

uti.create_boxplot(stat_values, timescales, selected, placeholders)

# Update existing selectors
previous_selection = uti.read_json_to_dict('selection.json')

# Set up widgets with previous settings
country_selector.value = previous_selection.get('country', placeholders['country'])
adm1_subarea_selector.value = previous_selection.get('adm1_subarea', placeholders['adm1_subarea'])
adm2_subarea_selector.value = previous_selection.get('adm2_subarea', placeholders['adm2_subarea'])
timescale_selector.value = previous_selection.get('timescale', placeholders['timescale'])
year_selector.value = previous_selection.get('year', placeholders['year'])
year_widgets_btn.on_click(on_button_clicked)

# Display widgets
display(country_selector, adm1_subarea_selector, adm2_subarea_selector, timescale_selector, year_selector, year_widgets_btn, output_area)

uti.display_data_details(active_btn, selected, subset_data[index])

Country:  Italy
ADM1 subarea:  Nord-Est
ADM2 subarea:  no adm2 subarea selected...
Year:  1998
Timescale:  3 months 

Time values in the subset: 12
Latitude values in the subset: 14
Longitude values in the subset: 20 

Data sample:  [[-9.99900000e+03 -9.99900000e+03 -9.99900000e+03 -9.99900000e+03
   6.93787587e-01]
 [-9.99900000e+03 -9.99900000e+03 -9.99900000e+03  7.10923387e-01
   6.71962754e-01]
 [ 7.81756359e-01  7.88598743e-01  6.59171139e-01  6.02512805e-01
   5.53822168e-01]
 [ 6.79895365e-01  6.37505279e-01  4.48464948e-01  3.38852353e-01
   2.47576329e-01]
 [ 4.91366561e-01  4.10632395e-01  2.86515031e-01  2.66845584e-01
   1.80545017e-01]]

processed_subset, change_summary = uti.process_datarray(subset_data[index])
print(processed_subset, '\n')
print('Change summary:')
for key, val in change_summary.items():
    print(key, val)

<xarray.DataArray 'SPEI3' (time: 12, lat: 14, lon: 20)>
dask.array<where, shape=(12, 14, 20), dtype=float64, chunksize=(1, 14, 20), chunktype=numpy.ndarray>
Coordinates:
  * time     (time) datetime64[ns] 1998-01-01T06:00:00 ... 1998-12-01T06:00:00
  * lon      (lon) float64 9.25 9.5 9.75 10.0 10.25 ... 13.25 13.5 13.75 14.0
  * lat      (lat) float64 43.75 44.0 44.25 44.5 44.75 ... 46.25 46.5 46.75 47.0
Attributes:
    long_name:  Standardized Drought Index (SPEI3)
    units:      - 

Change summary:
invalid_values_replaced 444
invalid_ratio 0.13214285714285715
duplicates_removed 0
cftime_conversions 0

stat_values = uti.compute_stats(processed_subset, full_stats=False)

uti.create_linechart(stat_values, timescales, selected, placeholders)

# Update existing selectors
previous_selection = uti.read_json_to_dict('selection.json')

# Set up widgets with previous settings
country_selector.value = previous_selection.get('country', placeholders['country'])
adm1_subarea_selector.value = previous_selection.get('adm1_subarea', placeholders['adm1_subarea'])
adm2_subarea_selector.value = previous_selection.get('adm2_subarea', placeholders['adm2_subarea'])
timescale_selector.value = previous_selection.get('timescale', placeholders['timescale'])
year_range_selector.value = previous_selection.get('year_range')
year_range_widgets_btn.on_click(on_button_clicked)

# Display widgets
display(country_selector, adm1_subarea_selector, adm2_subarea_selector, timescale_selector, year_range_selector, year_range_widgets_btn, output_area)

uti.display_data_details(active_btn, selected, subset_data[index])

Country:  Italy
ADM1 subarea:  Nord-Est
ADM2 subarea:  no adm2 subarea selected...
Year range:  ('2014', '2024')
Timescale:  3 months 

Time values in the subset: 124
Latitude values in the subset: 14
Longitude values in the subset: 20 

Data sample:  [[-9.99900000e+03 -9.99900000e+03 -9.99900000e+03 -9.99900000e+03
   9.25071009e-01]
 [-9.99900000e+03 -9.99900000e+03 -9.99900000e+03  1.02011684e+00
   1.01443264e+00]
 [ 1.21878881e+00  1.29934835e+00  1.21117076e+00  1.21989841e+00
   1.21718637e+00]
 [ 1.27909222e+00  1.27673982e+00  1.20574073e+00  1.14198875e+00
   1.07900212e+00]
 [ 1.20232192e+00  1.17460068e+00  1.07761040e+00  1.18870283e+00
   1.17344467e+00]]

processed_subset, change_summary = uti.process_datarray(subset_data[index])
print(processed_subset, '\n')
print('Change summary:')
for key, val in change_summary.items():
    print(key, val)

<xarray.DataArray 'SPEI3' (time: 122, lat: 14, lon: 20)>
dask.array<getitem, shape=(122, 14, 20), dtype=float64, chunksize=(1, 14, 20), chunktype=numpy.ndarray>
Coordinates:
  * time     (time) datetime64[ns] 2014-01-01T06:00:00 ... 2024-02-01T06:00:00
  * lon      (lon) float64 9.25 9.5 9.75 10.0 10.25 ... 13.25 13.5 13.75 14.0
  * lat      (lat) float64 43.75 44.0 44.25 44.5 44.75 ... 46.25 46.5 46.75 47.0
Attributes:
    long_name:  Standardized Drought Index (SPEI3)
    units:      - 

Change summary:
invalid_values_replaced 4588
invalid_ratio 0.13214285714285715
duplicates_removed 2
cftime_conversions 4

stat_values = uti.compute_stats(processed_subset, full_stats=False)

uti.create_stripechart(stat_values, timescales, selected, placeholders)

uti.create_stripechart(stat_values, timescales, selected, placeholders, 'year')

countries' boundaries source: www.geoboundaries.org